from django.core.management.base import BaseCommand
from django.contrib.auth.models import Group
from django.contrib.auth import get_user_model
User = get_user_model()


class Command(BaseCommand):
    help = 'Create a superuser and add to admin group'

    def handle(self, *args, **options):
        # 创建用户组
        admin_group, created = Group.objects.get_or_create(name='管理员')
        researcher_group, created = Group.objects.get_or_create(name='研究员')
        student_group, created = Group.objects.get_or_create(name='学生')
        # 接受用户输入
        user_specs = [
            {
                'username': '王院长',
                'password': '12345678',
                'mobile': '19999988888',
                'email': '',
                'group': admin_group
            },
            {
                'username': '张老师',
                'password': '12345678',
                'mobile': '18888877777',
                'email': '',
                'group': researcher_group
            },
            {
                'username': '张三丰',
                'password': '12345678',
                'mobile': '17777766666',
                'email': '',
                'group': student_group
            }
        ]
        # 创建超级用户
        for us in user_specs:
            username = us['username']
            password = us['password']
            mobile = us['mobile']
            email = us['email']
            group = us['group']

            # 检查用户是否存在，如果不存在，就创建
            try:
                user = User.objects.get(username=username)
            except:
                if group == admin_group:
                    user = User.objects.create_superuser(
                        username=username,
                        mobile=mobile,
                        email=email,
                        password=password
                    )
                else:
                    user = User.objects.create_user(
                        username=username,
                        mobile=mobile,
                        email=email,
                        password=password
                    )
                user.groups.add(group)

        self.stdout.write(self.style.SUCCESS('Superuser created successfully!'))
